// usbhotkey: Turn USB keyboard events into X11 keyboard events.
// Copyright (C) 2007 Lars Krueger
//
// usbremap: 2012 windexlight, renamed from usbhotkey
//
// This program is free software; you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation; either version 2 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program; if not, write to the Free Software
// Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA

// Portions of this code (overlay image) are ported from xteddy-2.0.1. 
// The ruby specific code is taken from various tutorials and forum entries
// from all over the net.

#include <stdlib.h>
#include <signal.h>
#include <sys/time.h>
#include <config.h>

// Undef the symbols that may collide with ruby
#undef PACKAGE_NAME
#undef PACKAGE_VERSION
#undef PACKAGE_STRING
#undef PACKAGE_TARNAME

#include <hid.h>
#include <ruby.h>
#include <X11/Xlib.h>
#include <X11/extensions/XTest.h>
#include <X11/XKBlib.h>

#include "constants.h"

//#define DEBUG_PACKET
#ifndef unk
#define unk 0
#endif

// This table maps USB HID scancodes (the type that will be returned by
// USB keyboards this utility interfaces with) to AT type scancodes.
// Adding 8 to these values will make them usable for looking
// up keysyms with Xkb.  This is pulled verbatim (just renamed) from
// the linux kernel source: hid-input.c
static const unsigned char usb_keymap[256] = {
	  0,  0,  0,  0, 30, 48, 46, 32, 18, 33, 34, 35, 23, 36, 37, 38,
	 50, 49, 24, 25, 16, 19, 31, 20, 22, 47, 17, 45, 21, 44,  2,  3,
	  4,  5,  6,  7,  8,  9, 10, 11, 28,  1, 14, 15, 57, 12, 13, 26,
	 27, 43, 43, 39, 40, 41, 51, 52, 53, 58, 59, 60, 61, 62, 63, 64,
	 65, 66, 67, 68, 87, 88, 99, 70,119,110,102,104,111,107,109,106,
	105,108,103, 69, 98, 55, 74, 78, 96, 79, 80, 81, 75, 76, 77, 71,
	 72, 73, 82, 83, 86,127,116,117,183,184,185,186,187,188,189,190,
	191,192,193,194,134,138,130,132,128,129,131,137,133,135,136,113,
	115,114,unk,unk,unk,121,unk, 89, 93,124, 92, 94, 95,unk,unk,unk,
	122,123, 90, 91, 85,unk,unk,unk,unk,unk,unk,unk,111,unk,unk,unk,
	unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,
	unk,unk,unk,unk,unk,unk,179,180,unk,unk,unk,unk,unk,unk,unk,unk,
	unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,unk,
	unk,unk,unk,unk,unk,unk,unk,unk,111,unk,unk,unk,unk,unk,unk,unk,
	 29, 42, 56,125, 97, 54,100,126,164,166,165,163,161,115,114,113,
	150,158,159,128,136,177,178,176,142,152,173,140,unk,unk,unk,unk
};


//################################# VARIABLES ################################
//=========================== HID DEVICE INTERFACE ===========================

static HIDInterface* theHid;           // Handle for opened Human Interface Device
static unsigned theHidIsOpen = 0;      // Set non-zero if ruby script called the "open device" function
static unsigned thePrgShouldQuit = 0;  // Set by signal handler to quit the program.
static unsigned theEndPoint;           // Number of the device endpoint.


//=========================== USB KEYBOARD PACKETS ===========================
// We hardcode the packet length here, although it is theoretically possible
// that a USB keyboard sends other packets.
// 
// Each USB packet consists of a byte indicating the state of the shift, ctrl,
// alt, and windows key. The next byte is reserved.
// Then 6 bytes follow that hold the the ids of the currently pressed keys.
// Such a packet is sent every time a key is pressed or released.

#define thePacketLen 6                           // Number of usable bytes for keycode
static char thePktBuf[ 2 ][ 2+thePacketLen ];    // Storage space for 2 packets: the last one seen and that of the current event
static char* thePacket = thePktBuf[ 0 ];         // Points to current packet
static char* theLastPacket = thePktBuf[ 1 ];     // Points to packet of last event
static char* theTrueLastPacket = thePktBuf[ 1 ]; // Points to the last received packet (persistent across idle events)
static struct timeval theLastTime;               // The time of the last event.


//====================== RUBY VARIABLES AND CALLBACK IDS =====================
// An table to keep track of information for each ruby method we want to call from C

enum
{
  RFUNC_KEYDN,  // KEYDN and KEYUP need to be the 0 and 1 because they serve
  RFUNC_KEYUP,  // double duty in the key event processing functions
  RFUNC_KEYEV,
  RFUNC_MARK,
  NUM_RUBY_FUNCS
};

typedef struct
{
  const char* name;
  ID func;
  unsigned def;
}
rubyFunc;

static rubyFunc rubyFuncs[] =
{
  { "UHK_USB_keydn", 0, 0 },
  { "UHK_USB_keyup", 0, 0 },
  { "UHK_USB_keyev", 0, 0 },
  { "UHK_USB_mark",  0, 0 }
};


//==================================== X11 ===================================

static Display* the_display = NULL;   // X11 server connection
#define INTERKEY_DELAY 0              // Number of milliseconds to wait after a keypress was sent.


//############################## INFRASTRUCTURE ##############################

// Displays error message and help.
static void errorHelp( const char* msg )
{
  fprintf( stderr, "usbhotkey: %s\n", msg );
  fputs( 
      "usbhotkey: Intercept USB keyboard events and send them to X.\n"
      "USAGE: usbhotkey [OPTIONS] <script>\n"
      "\n"
      "Options:\n"
      "  -h\n"
      "  --help -- Display help.\n"
      "\n"
      "Parameters:\n"
      "  <script> -- Ruby script containing the code to run on keypress.\n"
      , stderr );
  exit( EXIT_FAILURE );
}

// Signal handler, set a flag to quit.
static void quitMe( int sigNr )
{
  thePrgShouldQuit = 1;
}

// Check for a ruby exception and print some help.
static void handle_ruby_error( int status )
{
  if( status )
  {
    unsigned c;
    VALUE lasterr = rb_gv_get( "$!" );
    VALUE message = rb_obj_as_string( lasterr );
    VALUE ary = rb_funcall( rb_errinfo(), rb_intern("backtrace"), 0 );
    if ( !strcmp( RSTRING_PTR(message), "exit" ) )
    {
      quitMe( 1 );
      return;
    }
    printf( "usbhotkey: %s\n", RSTRING_PTR(message) );
    printf( "usbhotkey: Backtrace:\n" );
    for( c=0; c<RARRAY_LEN(ary); c++ ) 
      printf( "usbhotkey: \tfrom %s\n", RSTRING_PTR( RARRAY_PTR(ary)[c] ) );
    errorHelp( "Error executing start script\n" );
  }
}

// Call a ruby method from C - do nothing if not defined
static void callRubyFunc( unsigned func, int n, ... )
{
  if ( rubyFuncs[func].def )
  {
    VALUE *argv;
    va_list ar;
    if ( n > 0 )
    {
      long i;
      va_start( ar, n );
      argv = ALLOCA_N( VALUE, n );
      for ( i=0; i < n; i++ )
      {
          argv[i] = va_arg( ar, VALUE );
      }
      va_end( ar );
    }
    else
    {
      argv = 0;
    }
    rb_funcall2( rb_mKernel, rubyFuncs[func].func, n, argv );
  }
}


//############################ HID DEVICE HANDLING ###########################

// Ruby callback to open a device given its vendor and product id
static VALUE connectHid( VALUE self, VALUE vendorIdVal, VALUE productIdVal, VALUE endPointVal )
{
  if( theHidIsOpen )
    rb_raise( rb_eRuntimeError, "HID already connected." );

  unsigned vendorId = NUM2UINT( vendorIdVal );
  unsigned productId = NUM2UINT( productIdVal );
  theEndPoint = NUM2UINT( endPointVal );

  HIDInterfaceMatcher matcher = { vendorId, productId, NULL, NULL, 0 };

  printf( "usbhotkey: Connecting to %04x:%04x\n", vendorId, productId );

  hid_return ret = hid_force_open( theHid, 0, &matcher, 3 );
  if( ret != HID_RET_SUCCESS )
    rb_raise( rb_eRuntimeError, "Can't connect HID." );

  theHidIsOpen = 1;
  return self;
}

// Closes HID device and X11 connection
static void closeAll()
{
  if( !theHid )
    return;
  
  if( theHidIsOpen )
  {
    // Reset the USB device before closing the HID handle
    // This will cause the device to be reclaimed by the kernel drivers
    // as if it was unplugged and replugged.  In short - give it back to the OS.
    usb_reset( theHid->dev_handle );
    hid_close( theHid );
    theHidIsOpen = 0;
  }
  hid_delete_HIDInterface( &theHid );
  hid_cleanup();

  if( the_display )
  {
    XCloseDisplay( the_display );
    the_display = NULL;
  }
}


//########################### USB PACKET PROCESSING ##########################

// Check all keys in test that are not in set and call ruby functions for them
static void process_USB_keys( char* test, char* set, unsigned evtype )
{
  unsigned i, j;
  test += 2;
  set += 2;
  for ( i=0; i<thePacketLen; ++i )
  {
    if ( test[i] )
    {
      unsigned found = 0;
      for ( j=0; j<thePacketLen; ++j )
      {
        if ( test[i] == set[j] )
        {
          found = 1;
          break;
        }
      }
      if( !found )
      {
        // Call both the up/down function and the generic function
        // Generally only one will be defined in any given script
        VALUE key = INT2FIX( test[i] );
        callRubyFunc( evtype, 1, key );
        callRubyFunc( RFUNC_KEYEV, 2, key, INT2FIX( evtype ) );
      }
    }
  }
}

// Check all shift key bits in test that are not in set and call ruby functions for them
static void process_USB_shift_keys( unsigned char test, unsigned char set, unsigned evtype )
{
  unsigned char mask;
  int i;
  for ( i=0; i<8; ++i )
  {
    mask = 1 << i;
    if ( test & mask )
    {
      if ((set & mask) == 0 )
      {
        // Call both the up/down function and the generic function
        // Generally only one will be defined in any given script
        VALUE key = INT2FIX( 224 + i );
        callRubyFunc( evtype, 1, key );
        callRubyFunc( RFUNC_KEYEV, 2, key, INT2FIX( evtype ) );
      }
    }
  }
}

// Process the current packet. This function calls ruby functions which might
// raise an exception. In this case the remainder of the function is not executed.
static VALUE handlePacket( VALUE args )
{
  // Get the time since the last event
  struct timeval thisTime;
  int timeElapsed;
  gettimeofday( &thisTime, NULL );
  timeElapsed=1000*(thisTime.tv_sec-theLastTime.tv_sec)+(thisTime.tv_usec-theLastTime.tv_usec)/1000;

  // Process key releases
  callRubyFunc( RFUNC_MARK, 2, INT2FIX(MARK_START_RELEASE), INT2FIX(timeElapsed) );
  process_USB_keys( theLastPacket, thePacket, RFUNC_KEYUP );
  process_USB_shift_keys( theLastPacket[0], thePacket[0], RFUNC_KEYUP );
  
  // Process key presses
  callRubyFunc( RFUNC_MARK, 2, INT2FIX(MARK_END_RELEASE_START_PUSH), INT2FIX(timeElapsed) );
  process_USB_shift_keys( thePacket[0], theLastPacket[0], RFUNC_KEYDN );
  process_USB_keys( thePacket, theLastPacket, RFUNC_KEYDN );
  callRubyFunc( RFUNC_MARK, 2, INT2FIX(MARK_END_PUSH), INT2FIX(timeElapsed) );

  theLastTime = thisTime;
  return 0;
}


//############################ X11 KEY GENERATION ############################

// Send X11 key press event
static VALUE x11_keydn( VALUE self, VALUE keycode )
{
  unsigned kc = FIX2UINT( keycode );
  XTestFakeKeyEvent( the_display, XKeysymToKeycode( the_display, kc ), True, INTERKEY_DELAY );
  return 0;
}

// Send X11 key release event
static VALUE x11_keyup( VALUE self, VALUE keycode )
{
  unsigned kc = FIX2UINT( keycode );
  XTestFakeKeyEvent( the_display, XKeysymToKeycode( the_display, kc ), False, INTERKEY_DELAY );
  return 0;
}

// Send X11 key press/release event (depending on value of "up")
static VALUE x11_keyev( VALUE self, VALUE keycode, VALUE up )
{
  unsigned kc = FIX2UINT( keycode );
  Bool down = True;
  if ( FIX2UINT(up) > 0 )
    down = False;
  XTestFakeKeyEvent( the_display, XKeysymToKeycode( the_display, kc ), down, INTERKEY_DELAY );
  return 0;
}

// Send X11 button press event
static VALUE x11_btndn( VALUE self, VALUE button )
{
  XTestFakeButtonEvent( the_display, FIX2UINT( button ), True, INTERKEY_DELAY );
  return 0;
}

// Send X11 button release event
static VALUE x11_btnup( VALUE self, VALUE button )
{
  XTestFakeButtonEvent( the_display, FIX2UINT( button ), False, INTERKEY_DELAY );
  return 0;
}

// Send X11 button press/release event (depending on value of "up")
static VALUE x11_btnev( VALUE self, VALUE button, VALUE up )
{
  Bool down = True;
  if ( FIX2UINT(up) > 0 )
    down = False;
  XTestFakeButtonEvent( the_display, FIX2UINT(button), down, INTERKEY_DELAY );
  return 0;
}

// Get the KeySym associated with the given Keycode
static VALUE x11_get_keysym( VALUE self, VALUE code )
{
  // Take the raw USB HID scancode, translate it to a more traditional AT type 
  // scancode using the usb_keymap table, and further translate it to a usable
  // Xkb keycode by adding 8 to it (don't ask me why Xkb couldn't just use the
  // regular AT codes).  Finally, use the result to look up the associated keysym.
  return INT2FIX( XkbKeycodeToKeysym( the_display, usb_keymap[ FIX2INT(code) ] + 8, 0, 0 ) );
}

// Find out if a key is pressed or not
static VALUE uhk_is_key_pressed( VALUE self, VALUE code )
{
  char c = (char)FIX2INT( code );
  char *p = theTrueLastPacket + 2;
  unsigned i;
  for ( i=0; i<thePacketLen; ++i )
  {
    if ( p[ i ] == c )
    {
      return INT2FIX( 1 );
    }
  }
  return INT2FIX( 0 );
}


//################################### RUBY ###################################

// Initialize array of ruby methods to be called from C
static void initRubyFuncs( void )
{
  // Set up ruby methods to be called from C
  long i;
  for ( i=0; i<NUM_RUBY_FUNCS; i++ )
  {
    rubyFuncs[i].func = rb_intern( rubyFuncs[i].name );
    rubyFuncs[i].def = rb_obj_respond_to( rb_mKernel, rubyFuncs[i].func, 1 );
  }
}


//################################### MAIN ###################################

int main( int argc, char** argv )
{
  int status = 0;
  char* script;
  hid_return ret;

  // Process commandline
  int opti = 1;
  while ( opti < argc )
  {
    if( !strcmp( argv[opti], "-h" ) || !strcmp( argv[opti], "--help" ) )
      errorHelp( "Display help text." );
    else
      break;
  }

  if ( opti >= argc )
    errorHelp( "No script given." );
  
  script = argv[opti];

  // Init libhid
  hid_init();
  hid_set_debug( HID_DEBUG_NONE );
  hid_set_debug_stream( stderr );
  hid_set_usb_debug( 1000 );

  theHid = hid_new_HIDInterface();
  atexit( closeAll );
  signal( SIGINT, quitMe );
  signal( SIGTERM, quitMe );

  memset( theLastPacket, 0, thePacketLen+2 );

  // Init X11 connection and Test extension
  // Explicity pass the display (instead of using NULL) so that the call will
  // succeed when being invoked from a udev rule (where the DISPLAY env variable doesn't
  // seem to be set).  Can't say for sure at this point if there's any disadvantage
  // in any situation to using an explicit display like this.
  the_display = XOpenDisplay( ":0.0" );
  if ( !the_display )
    errorHelp( "Can't open X11 connection." );

  // Init ruby
  ruby_init();
  ruby_init_loadpath();

  // Set up C functions to be called from ruby
  rb_define_global_function( "UHK_connect", connectHid, 3 );
  rb_define_global_function( "UHK_keydn", x11_keydn, 1 );
  rb_define_global_function( "UHK_keyup", x11_keyup, 1 );
  rb_define_global_function( "UHK_btndn", x11_btndn, 1 );
  rb_define_global_function( "UHK_btnup", x11_btnup, 1 );
  rb_define_global_function( "UHK_keyev", x11_keyev, 2 );
  rb_define_global_function( "UHK_btnev", x11_btnev, 2 );
  rb_define_global_function( "UHK_get_keysym", x11_get_keysym, 1 );
  rb_define_global_function( "UHK_is_key_pressed", uhk_is_key_pressed, 1 );
    
  create_mark_constants();
  create_usb_constants();
  create_win_constants();

  // Load ruby script
  printf( "usbhotkey: Loading script %s\n", script );
  rb_load_protect( rb_str_new2( script ), 0, &status );
  handle_ruby_error( status );
  
  // Init ruby methods to be called from C
  initRubyFuncs();
  
  // Perform USB and X11 event processing
  if ( theHidIsOpen )
  {
    printf( "usbhotkey: Starting event loop\n" );

    gettimeofday( &theLastTime, NULL );
    theLastTime.tv_sec -= 10;

    // Event loop
    while ( !thePrgShouldQuit )
    {
      memset( thePacket, 0, thePacketLen+2 );

      // Read a packet
      ret = hid_interrupt_read( theHid, theEndPoint, thePacket, thePacketLen+2, 100 );
      
      if( ret == HID_RET_SUCCESS )
      {
        theTrueLastPacket = thePacket;
        XTestGrabControl( the_display, True );
        
#ifdef DEBUG_PACKET
        unsigned i;
        printf( "usbhotkey: Packet received: ");
        for( i=0; i<thePacketLen+2; ++i)
          printf( "%02x ", thePacket[i]);
        printf( "\n");
#endif

        // Find set differences and call ruby functions
        rb_protect( handlePacket, 0, &status );
        handle_ruby_error( status );

        // Switch roles of packet and last packet
        {
          char* h = thePacket;
          thePacket = theLastPacket;
          theLastPacket = h;
        }
        XTestGrabControl( the_display, False );
        XSync( the_display, False );
      }
      else
#ifdef LIBHID_HAS_TIMEOUT
      if( ret == HID_RET_TIMEOUT )
#endif
      {
        // Call the mark function indicating an idle event
        {
          // Get the time since the last event
          struct timeval thisTime;
          int timeElapsed;
          gettimeofday( &thisTime, NULL );
          timeElapsed = 1000*(thisTime.tv_sec-theLastTime.tv_sec)+(thisTime.tv_usec-theLastTime.tv_usec)/1000;
          XTestGrabControl( the_display, True );
          callRubyFunc( RFUNC_MARK, 2, INT2FIX(MARK_IDLE), INT2FIX(timeElapsed) );
          XTestGrabControl( the_display, False );
          XSync( the_display, False );
          theLastTime = thisTime;
        }
      }
#ifdef LIBHID_HAS_TIMEOUT
      else
        break;
#endif
    }
  }
  else
  {
    printf( "usbhotkey: No device opened by script. Exiting.\n" );
  }

  closeAll();
  return EXIT_SUCCESS;
}
